#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
PDF批量加水印工具
将同一文件夹中的所有PDF文件添加水印，并将结果保存到新文件夹中
"""

import os
import sys
import argparse
from PyPDF2 import PdfReader, PdfWriter
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def get_watermark(watermark_path):
    """
    获取水印文件
    
    Args:
        watermark_path: 水印PDF文件路径
    
    Returns:
        水印PDF的第一页，如果失败则返回None
    """
    try:
        watermark_pdf = PdfReader(watermark_path)
        return watermark_pdf.pages[0]  # 返回水印PDF的第一页
    except Exception as e:
        logger.error(f"无法读取水印文件: {e}")
        return None

def add_watermark_to_pdf(input_pdf_path, watermark_page, output_pdf_path):
    """
    给单个PDF文件添加水印
    
    Args:
        input_pdf_path: 输入PDF文件路径
        watermark_page: 水印页面对象
        output_pdf_path: 输出PDF文件路径
    
    Returns:
        成功返回True，失败返回False
    """
    try:
        # 读取原始PDF
        input_pdf = PdfReader(input_pdf_path)
        output_pdf = PdfWriter()
        
        # 为每一页添加水印
        for i in range(len(input_pdf.pages)):
            page = input_pdf.pages[i]
            page.merge_page(watermark_page)  # 将水印合并到原始页面
            output_pdf.add_page(page)
        
        # 保存新PDF
        with open(output_pdf_path, 'wb') as output_file:
            output_pdf.write(output_file)
        
        logger.info(f"已成功添加水印: {os.path.basename(input_pdf_path)}")
        return True
    except Exception as e:
        logger.error(f"添加水印失败 {os.path.basename(input_pdf_path)}: {e}")
        return False

def process_folder(input_folder, output_folder=None, watermark_filename="watermark.pdf"):
    """
    处理文件夹中的所有PDF文件并添加水印
    
    Args:
        input_folder: 输入文件夹路径
        output_folder: 输出文件夹路径，如果为None则自动创建
        watermark_filename: 水印文件名称
    
    Returns:
        成功处理的文件数量
    """
    # 确保输入文件夹存在
    if not os.path.exists(input_folder):
        logger.error(f"输入文件夹不存在: {input_folder}")
        return 0
    
    # 水印文件路径
    watermark_path = os.path.join(input_folder, watermark_filename)
    if not os.path.exists(watermark_path):
        logger.error(f"水印文件不存在: {watermark_path}")
        return 0
    
    # 获取水印页面
    watermark_page = get_watermark(watermark_path)
    if watermark_page is None:
        logger.error("处理水印文件时出错")
        return 0
    
    # 设定输出文件夹
    if output_folder is None:
        output_folder = os.path.join(input_folder, "水印处理结果")
    
    # 创建输出文件夹（如果不存在）
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
        logger.info(f"已创建输出文件夹: {output_folder}")
    
    # 计数器
    successful_count = 0
    
    # 处理所有PDF文件
    for filename in os.listdir(input_folder):
        if filename.lower().endswith('.pdf') and filename != watermark_filename:
            input_path = os.path.join(input_folder, filename)
            output_path = os.path.join(output_folder, filename)
            
            if add_watermark_to_pdf(input_path, watermark_page, output_path):
                successful_count += 1
    
    logger.info(f"批处理完成: 共处理 {successful_count} 个文件")
    return successful_count

def main():
    """
    主函数 - 解析命令行参数并执行水印处理
    """
    parser = argparse.ArgumentParser(description="PDF批量加水印工具")
    parser.add_argument("-i", "--input", help="输入文件夹路径", required=True)
    parser.add_argument("-o", "--output", help="输出文件夹路径（可选）")
    parser.add_argument("-w", "--watermark", help="水印文件名（默认为'watermark.pdf'）", default="watermark.pdf")
    
    args = parser.parse_args()
    
    print("=" * 50)
    print("PDF批量加水印工具")
    print("=" * 50)
    print(f"输入文件夹: {args.input}")
    print(f"输出文件夹: {args.output if args.output else '默认（输入文件夹/水印处理结果）'}")
    print(f"水印文件: {args.watermark}")
    print("=" * 50)
    
    count = process_folder(args.input, args.output, args.watermark)
    
    if count > 0:
        print(f"\n处理成功! 已处理 {count} 个PDF文件")
    else:
        print("\n处理失败! 请检查错误信息")

if __name__ == "__main__":
    main() 